#!/usr/bin/env python
#from __future__ import print_function as fprint
import sys
import os
import hmmer3  #minor modified version of hmmer3
#import MySQLdb
import sqlite3
import numpy as np
import csv
import urllib
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import IUPAC
from optparse import OptionParser
from bio.parsers import fasta as fasta_parser
from bio import blosum
#import auto_pdz
#import croc

#12/11/2012
#Will score all the pdz domains from query sequence (if there is more than one)
#change from mysql to sqlite

#11/10 change vpeptide expression to work with situation that
#      peptide len is 5

#known issue:
#query peptide contain 'X' or not AA
#weightM[i][p] will return error

AminoAcids = list('ACDEFGHIKLMNPQRSTVWY')
DBUSER = 'root'
DBPASSWD = '19854mdd'
"""
Version: 0.9 1003
fix SQL query problem
!='.' -> not in '.-'

Version: 0.8 0918
Fix hmmsequence, targetsequence
performance increase
add roc

Version: 0.7 0818
Add score normalization into output 

#1. call hmmscan (try to find it from J program) os.system
#2. parse hmmscan file to locate SDRs
#3. for each SDR res, check blosum62 for positive score
#4. (now having the residues from step 3)
#   select () from SDR where blah blah
#    select residues from peptide that bind to the same residues from step 3
#5. build matrix 
#os.system()
"""

def hmmscan_call(hmmdb, fa, eval):
    cmd = 'hmmsearch -E '+str(eval)+' '+hmmdb+' '+fa+'> tmp.hmmscan'
    os.system(cmd)
    hmmscan_out = open('tmp.hmmscan', 'r')
    return hmmscan_out
    
'''
def comb(items, _rs = []):
    #take a list of lists as input, return all possible [combination?]
    #print _tt
    _ts=[]
    _pop = items.pop()
    if type(_pop) is list:
        for i in _pop:
            _ts = [i]+items           #change to also work for str?
            if type(_ts[-1]) is list:
                comb(_ts, _rs)
            else:
                _tt.append(_rs)
    return _rs
'''

def isSDRhit(q_sdrs, d_sdrs, submatrix="blosum62",threhold=0):
    #two dictionaries of sdrs
    #return True is they are similar
    _submatrix = blosum.Matrix(submatrix)
    _isHit = True
    for key in d_sdrs:
        
        #print q_sdrs[key], d_sdrs[key]
        if q_sdrs[key] not in '.-' and d_sdrs[key] not in '.-':
            if _submatrix(q_sdrs[key], d_sdrs[key]) <= threhold:
                _isHit = False
        else:
            _isHit = False
            #print _submatrix(q_sdrs[key], d_sdrs[key])
    #print _isHit
    #raw_input()
            
    return _isHit

#Codes copy from Jonathan's version
class PDZError(Exception):
    pass

class MotifError(Exception):
    pass

def similar_sdrs(c_sdrs, dbname, blosum_matrix):
    """Return list of domain_ids with similar SDRs."""
    M = blosum.Matrix(blosum_matrix)
    con = sqlite3.connect(dbname)
    cur = con.cursor()
    domain_ids = None

    for key in c_sdrs.iterkeys():
        #print "Key:", key
        if key.startswith('GLGF'):
            motif = 'GLGF'
            offset = int(key[4:])
        elif key.startswith('GD'):
            motif = 'GD'
            offset = int(key[2:])
        else:
            raise MotifError("Motif needs to be either GLGF or GD.")
        
        if c_sdrs[key] in AminoAcids:
            residues = M.substitutions(c_sdrs[key], 1)
        
            residues = ','.join(['"{0}"'.format(r) for r in residues])
            #print residues
            subquery = """\
            SELECT DISTINCT domainID
            FROM sdr_pdz
            WHERE motif = '{0}'
            AND offset = {1}
            AND residue IN ({2})""".format(motif, offset, residues)
            cur.execute(subquery)
            
                
            if domain_ids is None:
                domain_ids = set([row[0] for row in cur])
            else:
                domain_ids = domain_ids.intersection(set([row[0] for row in cur]))
    #print "domain_ids:", domain_ids
    cur.close()
    con.close()
    return domain_ids

def build_query(peppos, ids):
    """
    MySQL is hopeless with subqueries (they run very slowly); therefore,
    seperate queries are made and the results combined in Python.  Tests show
    this to be much quicker.
    """
    #ids = similar_sdrs(c_sdrs, dbname, blosum_matrix)
    # TO-DO: is peppos counting from C-terminal?  i.e., is peppos = 0
    # the very end?
    sql = """\
    SELECT substr(aa.seq, {0}, 1) AS ss, count(*) AS cc
    FROM (
          SELECT seq
          FROM peptide_pdz
          WHERE domainID IN ({1})
         ) AS aa
    GROUP BY ss"""
    position = 5 - peppos #Please confirm this is right?
    residues = ','.join(["'{0}'".format(x) for x in ids])
    if len(ids) == 0:
        raise PDZError('intersection is empty')
    sql = sql.format(position, residues)
    #print ('SQL', sql)
    return sql
#-------------------------------------------------------------------------------

    

def sql_q(glGf_off, GD_off, i_sequence, peppos, blosumM, j_sequence, dbname, masking = []):

    #return FreqMatrix of ONE position of the peptide
    #i_sequence is the target sequence, j_sequence is the hmm sequence
    
    db = sqlite3.connect(dbname)
    curs = db.cursor()
    curs2 = db.cursor()
    curs.execute("""SELECT distinct s.motif, s.offset
                    FROM sdr as s
                    WHERE s.pep_pos = (?)""", (peppos,))
    motifs_set=curs.fetchall()

    #look up SDR positions from database
    #print motifs_set
    #M = blosum.Matrix(blosumM); #should accept user input for the var
    #print peppos, motifs_set
    cutoff = 0
    c_sdrs = dict() #compare sdrs | identify SDRs from query sequence
    
    for motif in motifs_set:
        #print motif
        #Create a list of SDR from query sequence
        if motif[0] == "GD":
            try:
                count = abs(motif[1])
                pointer = GD_off
                
                while count != 0:
                    #if there are gaps in between do not count them
                    
                    if motif[1] >=0:
                        pointer += 1
                    else:
                        pointer -= 1
                        
                    if j_sequence[pointer] not in '.-':
                        count -= 1
                        
                sdr = i_sequence[pointer]
                #print sdr
            except IndexError, e:
                #print peppos, e
                #print "ERR: SDR not in hmmscan alignment", motif[0], motif[1]
                #if the SDR is not in the alignment
                #not possible to find matches, return ALL 0
                freq = [0] * 20
                return freq
                
            c_sdrs[motif[0]+str(motif[1])] = sdr
            #print c_sdrs
        if motif[0] == "GLGF":
            try:
                count = abs(motif[1])
                pointer = glGf_off
                while count != 0:
                    #if there are gaps in between do not count them
                    
                    if motif[1] >= 0:
                        pointer += 1
                    else:
                        pointer -=1
                    
                    if j_sequence[pointer] not in '-.':
                        count -= 1    
                    
                sdr = i_sequence[pointer]
                #print "sdr:", sdr
            except (IndexError):
                #print "ERR: SDR not in hmmscan alignment", motif[0], motif[1]
                #if the SDR is not in the alignment
                #not possible to find matches, return ALL 0
                freq = [0] * 20
                return freq
            
            c_sdrs[motif[0]+str(motif[1])] = sdr
    #print motif, c_sdrs
    #print "Identification of SDRs from query sequence Done."
    
    results = []
    #print "255:", c_sdrs
    #---------CODE FROM JJE---------#
    #Looking for similar residues from database peptides
    freq = [0] * 20
    #ids = domain IDs, which have similar SDRs to the query domain.
    
    #print c_sdrs
    ids_2 = similar_sdrs(c_sdrs, dbname, blosumM)
    #print '263:',ids_2
    #print ids_2
    ids = []
    #MASKING AT HERE, ids <-- check!
    for i in ids_2:
        if i not in masking:
           ids.append(i) 
    
    #print 'Line 269 ids:', ids
    #print c_sdrs
    try:
        sql = build_query(peppos, ids)
    except PDZError:
        # There were no domains in the intersection. (so the IN clause is empty
        # and therefore an SQL error).
        return freq

    con = sqlite3.connect(dbname)
    cur = con.cursor()
    try:
        cur.execute(sql)
        for row in cur:
            try:
                idx = AminoAcids.index(row[0])
            except ValueError:
                pass
            else:
                freq[idx] = int(row[1])
    finally:
        cur.close()
        con.close()
    #print 'Freq:', freq
    return freq
    
#COPY and EDIT from Jonathan's Predikin3, kinase.py
#TODO: Create backgroud freq from database?
def create_weights(frequencies):
    #frequencies = self.frequencies
    #priors = self._get_priors()
    '''
    priors = {'PDZ':   {'A': 0.069, 'C': 0.017, 'D': 0.054, 'E': 0.069,
                       'F': 0.034, 'G': 0.071, 'H': 0.025, 'I': 0.048,
                       'K': 0.051, 'L': 0.092, 'M': 0.021, 'N': 0.038,
                       'P': 0.066, 'Q': 0.045, 'R': 0.060, 'S': 0.086,
                       'T': 0.054, 'V': 0.064, 'W': 0.010, 'Y': 0.027},
              }
    prior = priors['PDZ'] #using 'PDZ' at the moment
    '''
    #weights = np.zeros([5, 20])
    weights = {}
    
    sqrt = np.sqrt
    log2 = np.log2

    for i, row in enumerate(frequencies):
        weights[i] = {}


        N = sum(row)
        #if an entire row is 0, prediction fail? OR just score 0 for that row?
        
        
        if N > 0:
            for j,freq in enumerate(row):
                #print j, freq
                #p = prior[AminoAcids[j]]
                #weight = log2(((freq + sqrt(N / 20.0)) / (N + sqrt(N))) / p)
                #New style scoring matrix (Jonathan's Predikin paper)
                weight = (freq + sqrt(N / 20.0)) / (N + sqrt(N))
                weights[i][AminoAcids[j]] = weight
        else:
            #print "N = 0"
            return False
            raise 
        '''
        
        #This case, we still want to score even if one of more rows are all 0s.
        for j,freq in enumerate(row):
            #print j, freq
            p = prior[AminoAcids[j]]
            if N > 0:
                weight = log2(((freq + (sqrt(N) / 20.0)) / (N + sqrt(N))) / p)
                weights[i][AminoAcids[j]] = weight
            else:
                weights[i][AminoAcids[j]] = prior[AminoAcids[j]]
        '''
    #print weights
    return weights

def is_valid(M):
    # 0=invalid, 1=suspicious, 2=valid
    r = 'VALID'
    for row in M:
        if sum(row) == 0:
            r = 'INVALID'
            break
        elif sum(row) < 5:
            r = 'SUSPICIOUS'
            break
        
    if r == 'SUSPICIOUS' and sum(M[0]) > 5 and sum(M[2]) > 5:
        r = r + ' *'
    return r

def score_peptide(peptide, weightM, roc_scores = [], isbind = -1):
    score = 0
    minsum = 0
    maxsum = 0
    #print weightM
    for i, p in enumerate(peptide[::-1][:5]):
        skip = False
        #print peptide
        #print peptide[::-1][:5]
        #print i, p
        #if len(peptide) != 5:
        #    raw_input()
        try:
            raw_score = weightM[i][p]
        except:
            skip = True
            #raw_score = float(sum(weightM[i].itervalues())/ len(weightM[i]))
            #print "Len:", len(weightM[i]), raw_score
        #normalize the score
        #print raw_score
        if skip is False:
            minsum += min(weightM[i].itervalues())
            
            maxsum += max(weightM[i].itervalues())
            score += raw_score
        
    
    nscore = (score-minsum)/(maxsum-minsum)
    if nscore < 0:
        print "minsum:", minsum
        print "maxsum:", maxsum
        print "score:", score
        print peptide[::-1][:5]
    #print '|',nscore
    return nscore

    '''
    if isbind == -1:
        return nscore
    else:
        print nscore, isbind
        roc_scores.add(nscore, isbind)
    
        return roc_scores
    '''

def data2file(scores, dbname):
    i = dbname[-1]
    outfile = 'datafile{0}.dat'.format(i)
    data = scores.score_labels
    with open (outfile, 'w') as fout:
        for k in data:
            values = data[k]
            for v in values:
                writer = csv.writer(fout, delimiter = ' ', lineterminator = '\n')
                writer.writerow([k, v])




def getPWM(query_seq, dbname, masking = []):
    AminoAcids = list('ACDEFGHIKLMNPQRSTVWY')
    tmpPath = "scanHmmer.tmp"
    tmpFile = open(tmpPath, "w" )
        #####################################################################
    #   variables
    ##############
    #import csv table for validation

    hmmscan_thres = 0.01
    hmmdb = 'strucaln.hmm'
    options_blosum = "blosum62"
    options_length = 5
    #####################################################################
    previous_protein = ''
    previous_pdznum = ''
    
    
    
    #pdz_record = SeqIO.parse(infile, "fasta").next()

    if type(query_seq) == SeqRecord:
        SeqIO.write(query_seq, tmpFile, "fasta")
        tmpFile.close()
    else:
        print "Input Sequence Type Error."
        raise 
    read_data = hmmscan_call(hmmdb, tmpPath, hmmscan_thres)
    #raise
    #cnt_record = 0
    #cnt_hit = 0
    #cnt_domain = 0
    pwms = dict()
    for record in hmmer3.parse(read_data):
        #cnt_record += 1
        for hit in record:
            #cnt_hit += 1
            for domain in hit:
                #cnt_domain += 1
                
                glGf_off = domain.hmm_sequence.find('glGf')
                if glGf_off == -1:
                    print "GLGF motif not found in the HMM alignment."
                    return False
                else:
                    glGf_off = glGf_off + 3
                #index of glGf motif
                
                GD_off = domain.hmm_sequence.find('GD')
                if GD_off == -1:
                    print "GD motif not found in the Hmm alignment."
                    return False
                else:
                    GD_off = GD_off +1
                #index of GD motif
                #load the fasta file into fasta parser
                freqM = []   #output frequency matrix
                weightM = [] #output PWM matrix
                for p_position in range(options_length):
                    #careful with hmm_sequence and target_sequence
                    #print hit.name, p_position
                    freqM.append(sql_q(glGf_off, GD_off ,
                                       domain.target_sequence,
                                       p_position,
                                       options_blosum,
                                       domain.hmm_sequence,
                                       dbname,
                                       masking))
                    #print freqM
                    #print domain.target_sequence, domain.hmm_sequence
                weightM = create_weights(freqM)
                '''
                #=====OUTPUT=====
                #print weightM
                #print "Weight:", weightM
                if weightM:
                    #print "weight matrix ok!"
                    #x = is_valid(freqM)
                    score = 0
                    #print "len pep:", len(vpeptide)
                    if len(vpeptide) >= 4:
                        roc_scores = score_peptide(vpeptide, weightM, roc_scores, isbind)
                '''
                #print freqM
                if weightM is not False:
                    pwms[(domain.ali_from, domain.ali_to)] = weightM
                else:
                    print >> sys.stderr, 'weight matrix failed.'
                    #raise
                    pwms[(domain.ali_from, domain.ali_to)] = weightM
            try:
                os.unlink( tmpPath )
            except OSError:
                pass
            
            #print roc_scores
            #print '\t'.join([record.query, hit.name, domain.id, x])
            
    #data2file(roc_scores, dbname)
            #return weightM
    #print cnt_record, cnt_hit, cnt_domain
    if len(pwms) > 0:
        #index: (domain.ali_from, domain.ali_to)
        #value: weight matrix
        return pwms
    else:
        print "Length of PWM = 0", pwms
        return False

def seq2seqrecord(seq):
    #convert query sequence to SeqRecord format
    record = SeqRecord(Seq(seq,IUPAC.protein), id = "Query Sequence", name="Query",
                       description = "Query Sequence")
    
    return record

def fasta2seqrecord(infile):
    #convert fasta file to SeqRecord format
    
    record = SeqIO.parse(infile, "fasta").next()
    
    return record

def scoring(query_sequence, peptides, dbname, masking = []):
    pwms = getPWM(query_sequence, dbname, masking)

    results = dict()
    if pwms is not False:
        for k in pwms:
            
            scores = []
            #print pwms
            if pwms[k] is not False:
                for peptide in peptides:
                    if len(peptide) >= 5:
                        scores.append(score_peptide(peptide[0:5], pwms[k]))
                    else:
                        print peptide, "Peptide length is less than 5"
                        scores.append(False)
            else:
                scores = False
            results[k] = scores
        return results, len(pwms)
    else:
        return False

if __name__ == '__main__':
    #AminoAcids = list('ACDEFGHIKLMNPQRSTVWY')

    seq = SeqIO.parse("test.fasta", "fasta").next()
    peptides = ["AQTRL","AAAAA"]
    
    result = scoring (seq, peptides, "pdzvalid.db")
    
    if result is not False:
        scores, pdz_count = result
    
    
        print pdz_count, scores
    '''
    validtable = 'dummy_test.dat'
                        #a, b = main(validtable, dbname);
    for i in xrange(10):
        #testfile = 
        dbname = 'pdzvalid_{0}'.format(i + 1)
        roc, auc = main(validtable, dbname)
    '''
    #test(','/database/pdzvalid', [])
    '''
    for dbnum in range(0, 10):
        dbname = 'pdzvalid_' + str(dbnum + 1)
        validtable = 'test_group'+str(dbnum)+'.dat'
        a, b = main(validtable, dbname);
        print a 
        print b
    '''

